/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file extension_amp_update_scale.cpp
 */
#include <torch/csrc/autograd/custom_function.h>
#include <torch/extension.h>

#include "pytorch_npu_helper.hpp"
using torch::autograd::AutogradContext;
using torch::autograd::Function;
using tensor_list = std::vector<at::Tensor>;
using namespace at;

// register forward implementation for NPU device
// std::vector<at::Tensor>
at::Tensor my_op_impl_autograd(
    const at::Tensor& current_scale,
    const at::Tensor& growth_tracker,
    const at::Tensor& found_inf,
    float growth_factor,
    float backoff_factor,
    int32_t growth_interval)
{
    // alloc output memory
    at::Tensor updatedScaleOut = at::scalar_tensor(0.0);
    at::Tensor updatedGrowthTrackerOut = at::scalar_tensor(0);
    // std::vector<at::Tensor> result = {updatedScaleOut, updatedGrowthTrackerOut};

    // call aclnn interface to perform the computation
    EXEC_NPU_CMD(aclnnAmpUpdateScale, current_scale, growth_tracker, found_inf,
                 growth_factor, backoff_factor, growth_interval,
                 updatedScaleOut, updatedGrowthTrackerOut);

    return updatedScaleOut;
}

// bind the C++ interface to the Python interface using pybind
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("amp_update_scale_jit", &my_op_impl_autograd, "JIT version of AmpUpdateScale operator"); // 函数名，函数地址，注释
}
